import os,time, sys, math
import tensorflow as tf
import argparse
import numpy as np
from PIL import Image
from dataset import get_dataset
from utils import utils, helpers
from builders import model_builder

parser = argparse.ArgumentParser()
#/home/ubuntu/NN/checkpoints/latest_model_DeepLabV3_CamVid.ckpt-40000.data-00000-of-00001
#c:\Users\lamed\Desktop\NN\checkpoints\Encoder-Decoder-Skip\latest_model_Encoder-Decoder-Skip_CamVid.ckpt-0.data-00000-of-00001
parser.add_argument('--image', type=str, default=None, help='The image you want to predict on. ')
parser.add_argument('--checkpoint_path', type=str, default="/home/ubuntu/NN/checkpoints/latest_model_DeepLabV3_CamVid.ckpt-40000", help='The path to the latest checkpoint weights for your model.')
parser.add_argument('--crop_height', type=int, default=512, help='Height of cropped input image to network')
parser.add_argument('--crop_width', type=int, default=512, help='Width of cropped input image to network')
parser.add_argument('--model', type=str, default='DeepLabV3', help='The model you are using')
#parser.add_argument('--dataset', type=str, default="CamVid", required=False, help='The dataset you are using')
parser.add_argument('--frontend', type=str, default='ResNet101', help='The frontend you are using. See frontend_builder.py for supported models')
args = parser.parse_args()
label_names = ["Ring","section","leakwater"]
num_classes = len(label_names)
train_dir = "/home/ubuntu/桌面/segmentation/test_data.tfrecords"
# Initializing network
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess=tf.Session(config=config)
train_element,train_iter = get_dataset(train_dir,batch_size=1,is_training=False,reused=False)
#print(train_element)
net_input,net_output = train_element
#net_output = tf.placeholder(tf.float32,shape=[None,None,None,num_classes]) 
def Scandata_Inference(net_input,is_training = True):
    num_classes = len(label_names)
    network = model_builder.build_model(model_name=args.model,\
         frontend=args.frontend, net_input=net_input, \
             num_classes=num_classes, crop_width=args.crop_width, \
                 crop_height=args.crop_height, is_training=is_training)
    #network =tf.sigmoid(network)
    return network
#train_element = tf.reshape(train_element,[1,512,512,3])
network = Scandata_Inference(net_input,is_training=False)


sess.run(tf.global_variables_initializer())

print('Loading model checkpoint weights')
saver=tf.train.Saver(max_to_keep=0)
saver.restore(sess, args.checkpoint_path)
#c:\Users\lamed\Desktop\NN\checkpoints\latest_model_FC-DenseNet56_CamVid.ckpt-40000.data-00000-of-00001
#checkpoints/latest_model_DeepLabV3_CamVid.ckpt-40000.data-00000-of-00001
for step in range(10):
    mask,gt= sess.run([network,net_output])
    
    gt = np.reshape(gt,(512,512,3))*255
    mask = np.where(mask<0.6, 0, 1)
    mask = np.reshape(mask,(512,512,3))*255
    mask = mask.astype(np.uint8)   
    for i in range(3):
        im = Image.fromarray(mask[:,:,i])
        gt_im = Image.fromarray(gt[:,:,i])
        if i == 0 and step==6:
            im.show()
            gt_im.show()
        #imagearray_show(test_masks[:,:,i])
    #"""
#print("Testing image " + args.image)

print("")
print("Finished!")

